'''
Created on Dec 13, 2010

@author: mkiyer
'''
import os
import logging
import argparse
import lxml.etree as etree
from datetime import datetime
import collections

import seqdb
from seqdb.model.meta import Session
from seqdb.query import query_sequence_data, Filter

def generate_lane_xml(parent, seqdata, root_dir, remote_ip):
    root = etree.SubElement(parent, "lane")
    root.set("name", seqdata.name)    
    output_dir = os.path.join(root_dir, seqdata.name)
    elem = etree.SubElement(root, "output_dir")
    elem.text = output_dir     
    if remote_ip is not None:
        elem.set("remote", "True")     
        elem.set("ip", remote_ip)
    else:
        elem.set("remote", "False")
    # sample parameters
    elem = etree.SubElement(root, "run_id")
    elem.text = seqdata.name
    elem = etree.SubElement(root, "lib_id")
    elem.text = seqdata.seqlibrary.name
    elem = etree.SubElement(root, "sample_id")
    elem.text = seqdata.seqlibrary.experiment.name
    elem = etree.SubElement(root, "species")
    species = seqdata.seqlibrary.experiment.sample.organism.getvattr("species")    
    elem.text = species
    elem = etree.SubElement(root, "description")
    elem.text = seqdata.seqlibrary.experiment.sample.description
    # sequence data parameters
    elem = etree.SubElement(root, "center")
    elem.text = seqdata.instrument_run.getvattr("institution")
    elem = etree.SubElement(root, "date")
    elem.text = datetime.strftime(seqdata.instrument_run.run_date, '%Y-%m-%d')
    elem = etree.SubElement(root, "platform")
    elem.text = seqdata.instrument_run.getvattr('platform')
    elem = etree.SubElement(root, "read_length")
    elem.text = str(seqdata.getvattr('read_length'))
    elem = etree.SubElement(root, "fragment_layout")
    elem.text = seqdata.instrument_run.getvattr("fragment_layout")
    elem = etree.SubElement(root, "fragment_length_mean")
    elem.text = str(seqdata.seqlibrary.fragment_length_mean)
    elem = etree.SubElement(root, "quality_scores")
    elem.text = seqdata.instrument_run.getvattr("quality_scores")
    # sequence data files
    elem = etree.SubElement(root, "fastq_files")
    for mate in xrange(len(seqdata.files)):
        filename = seqdata.files[mate]
        file_elem = etree.SubElement(elem, "file", mate=str(mate))
        file_elem.text = filename


def generate_sample_xml(parent, name, seqdatalist, sample_dir, lane_dir, remote_ip):
    root = etree.SubElement(parent, "sample")
    root.set("name", name)
    # set output paths
    elem = etree.SubElement(root, "output_dir")
    elem.text = os.path.join(sample_dir, name)
    if remote_ip is not None:
        elem.set("remote", "True")     
        elem.set("ip", remote_ip)
    else:
        elem.set("remote", "False")
    # set patient id 
    patient_ids = set(seqdata.seqlibrary.experiment.sample.organism.name
                      for seqdata in seqdatalist)
    if len(patient_ids) > 1:
        logging.error("Found multiple patient IDs for data from a single sample")
        assert False
    elem = etree.SubElement(root, "patient_id")
    elem.text = list(patient_ids)[0]
    # set category 
    categories = set(seqdata.seqlibrary.experiment.sample.getvattr("clinical_diagnosis")
                     for seqdata in seqdatalist)
    if len(categories) > 1:
        logging.error("Found multiple categories for data from a single sample")
        assert False
    elem = etree.SubElement(root, "category")
    elem.text = list(categories)[0]
    # add lanes
    for seqdata in seqdatalist:
        generate_lane_xml(root, seqdata, lane_dir, remote_ip)        


def generate_analysis_xml(analysis_name, sample_dict, root_dir, remote_ip=None):
    # setup paths
    if root_dir is None:
        root_dir = os.getcwd()
    analysis_dir = os.path.join(root_dir, "analysis", analysis_name)
    sample_dir = os.path.join(root_dir, "sample")
    lane_dir = os.path.join(root_dir, "lane")
    # create xml tree
    root = etree.Element("analysis")
    root.set("name", analysis_name)
    elem = etree.SubElement(root, "output_dir")
    elem.text = analysis_dir
    if remote_ip is not None:
        elem.set("remote", "True")     
        elem.set("ip", remote_ip)
    else:
        elem.set("remote", "False")
    # add sample elements
    for sample_name, seqdatalist in sample_dict.iteritems():
        generate_sample_xml(root, sample_name, seqdatalist, sample_dir, lane_dir, remote_ip)
    return root

def main():
    logging.basicConfig(level=logging.DEBUG)
    parser = argparse.ArgumentParser()
    parser.add_argument('--filter', action="append", dest="filters", default=[], nargs="+", help="classname.attribute <op> [value]")
    parser.add_argument('--seqdb', default=None)
    parser.add_argument('--remote-ip', dest="remote_ip", default=None)
    parser.add_argument('--output-dir', dest="output_dir", default=None)
    parser.add_argument('name', help="Name for this analysis")
    
    options = parser.parse_args()
    if options.output_dir is None:
        parser.error("must specify an output directory with --output-dir")
    if not os.path.isabs(options.output_dir):
        parser.error("must specify an absolute path with --output-dir")
    # query database for sequence data
    seqdb.connect(options.seqdb)
    # create filters
    filters = []
    for args in options.filters:
        filters.append(Filter.from_args(args))        
    # group sequence data by sample
    samples = collections.defaultdict(lambda: []) 
    for seqdata in query_sequence_data(filters):
        sample_name = seqdata.seqlibrary.experiment.sample.name
        samples[sample_name].append(seqdata)
    # create analysis XML
    xmlelem = generate_analysis_xml(options.name, samples, options.output_dir, options.remote_ip)
    xmlstring = etree.tostring(xmlelem, pretty_print=True)
    print xmlstring    
    seqdb.disconnect()

if __name__ == '__main__': main()